# -*- coding: utf-8 -*-
"""
Created on Fri Jun  1 10:22:33 2012

@author: nzhao
"""
import numpy as np



class chain1e():
    def __init__(self, chainLen=1024):
        self.chain_length = chainLen
        self.hamiltonian = 2.0*np.identity(self.chain_length) \
            - np.diag(np.ones(self.chain_length-1),1) \
            - np.diag(np.ones(self.chain_length-1), -1)
       
    def diagonalize(self):
        [u,v] = np.linalg.eigh(self.hamiltonian)
        idx = u.argsort()
        self.energy = u[idx]
        self.state = v[:, idx]

class block(chain1e):
    def __init__(self, chainLen=8, truncate=8):
        chain1e.__init__(self, chainLen=chainLen)
        
        self.site_num = chainLen
        
        self.tunelling = np.zeros( (self.chain_length, self.chain_length) )
        self.tunelling[-1, 0] = -1.0

        self.m = truncate        
        
        self.diagonalize()
        self.RG_energy = self.energy[:self.m]
        self.RG_state = self.state[:, :self.m]
        
        self.wavefunction = self.RG_state.T

        
    def block_renormalize(self):
        self.RG_hamiltonian = np.dot( self.RG_state.conj().T, \
            np.dot(self.hamiltonian, self.RG_state) )
        self.RG_tunelling = np.dot( self.RG_state.conj().T, \
            np.dot(self.tunelling, self.RG_state) )
        
        self.block_grow()
        
        self.diagonalize()
        self.RG_energy = self.energy[:self.m]
        self.RG_state = self.state[:, :self.m]
        
        self.wavefunction = np.dot(self.wavefunction_dbl.T, self.RG_state).T
    
    def block_grow(self):
        self.hamiltonian = \
            np.r_[ \
                np.c_[self.RG_hamiltonian, self.RG_tunelling],\
                np.c_[self.RG_tunelling.conj().T, self.RG_hamiltonian] ]
        zeroMat= np.zeros((self.m, self.m))
        self.tunelling = \
            np.r_[ \
                np.c_[zeroMat, zeroMat],\
                np.c_[self.RG_tunelling, zeroMat] ]
        
        zeroWave = np.zeros(self.site_num)
        wave1 = [np.array([x, zeroWave]).flatten() for x in self.wavefunction]
        wave2 = [np.array([zeroWave, x]).flatten() for x in self.wavefunction]
        self.wavefunction_dbl= (np.r_[wave1, wave2])
        
        self.site_num = self.site_num * 2
        
